﻿#include "filemanager.h"
#include <filesystem>
#include <fstream>
#include "../CipherUtils/cipherutils.h"
#include "../AppConfig/appconfig.h"
#include <iostream>
#include "../Debug/DEBUG.hpp"
#include <codecvt>

bool FileManager::checkPath(const std::string &path, bool create)
{
    std::wstring_convert<std::codecvt_utf8<wchar_t>> converter;
    std::wstring filePath = converter.from_bytes(path);
    // 检查文件是否存在
    if (!std::filesystem::exists(filePath)) {
        // if do not need to create
        if (!create) {
            return false;
        }
        std::ofstream file(filePath);
        if (file.is_open()) {
            file.close();
        } else {
            return false;
        }
    } else {
        return true;
    }
    return false;
}

void FileManager::init()
{
    // check file dir path
    if (!std::filesystem::exists(FILES_DIR)) {
        // not exist
        std::filesystem::create_directory(FILES_DIR);
    }
    // 检查图片文件夹
    if (!std::filesystem::exists(PIC_DIR)) {
        // not exist
        std::filesystem::create_directory(PIC_DIR);
    }
}

std::pair<std::shared_ptr<unsigned char[]>, int> FileManager::readFromDisk(const std::string &path, uint64_t blockOffset, uint64_t size, bool decrypted)
{
    std::wstring_convert<std::codecvt_utf8<wchar_t>> converter;
    std::wstring filePath = converter.from_bytes(path);
    std::ifstream ifs(filePath, std::ios::binary | std::ios::in);
    if (!ifs.is_open()) {
        return {nullptr, 0};
    }
    int offset = 0;
    if(decrypted) {
        offset = blockOffset * ENCRYPTED_FILE_BLOCK_SIZE;
    } else {
        offset = blockOffset * FILE_BLOCK_SIZE;
    }


    ifs.seekg(offset, std::ios::beg);

    std::shared_ptr<unsigned char[]> buffer(new unsigned char[size]);

    int readSize = 0;

    ifs.read((char*)buffer.get(), size);
    readSize = ifs.gcount();

    // if it doesn't need to decrypt
    if (!decrypted) {
        return {buffer, readSize};
    }

    CipherUtils& utils = CipherUtils::getInstance();
    AppConfig& config = AppConfig::getInstance();

    std::string AESKey = config.runtimeValue("LocalAESKey");
    std::string AESIV = config.runtimeValue("LocalAESIV");

    auto res = utils.AES_Decrypt(buffer.get(), readSize,
                      (unsigned char*)AESKey.c_str(), (unsigned char*)AESIV.c_str());
    ifs.close();
    return res;
}

void FileManager::writeToDisk(const std::string &path, uint64_t blockOffset, const char *file, uint64_t size, bool encrypted)
{
    std::wstring_convert<std::codecvt_utf8<wchar_t>> converter;
    std::wstring filePath = converter.from_bytes(path);
    std::ofstream ofs(filePath, std::ios::binary | std::ios::out);
    if (!ofs.is_open()) {
        return ;
    }
    int offset = 0;
    if (encrypted) {
        offset = blockOffset * ENCRYPTED_FILE_BLOCK_SIZE;
    } else {
        offset = blockOffset * FILE_BLOCK_SIZE;
    }

    ofs.seekp(offset, std::ios::beg);

    if (!encrypted) {
        ofs.write(file, size);
        return ;
    }

    CipherUtils& utils = CipherUtils::getInstance();
    AppConfig& config = AppConfig::getInstance();

    std::string AESKey = config.runtimeValue("LocalAESKey");
    std::string AESIV = config.runtimeValue("LocalAESIV");

    auto res = utils.AES_Encrypt(file, size,
                                 (unsigned char*)AESKey.c_str(), (unsigned char*)AESIV.c_str());

    ofs.write((char *)res.first.get(), res.second);
    ofs.close();
}

void FileManager::writeToDisk(std::shared_ptr<std::ofstream> ofs, uint64_t blockOffset, const char *file, uint64_t size, bool encrypted)
{
    int offset = 0;
    if (encrypted) {
        offset = blockOffset * ENCRYPTED_FILE_BLOCK_SIZE;
    } else {
        offset = blockOffset * FILE_BLOCK_SIZE;
    }

    ofs->seekp(offset, std::ios::beg);
    if (!encrypted) {
        ofs->write(file, size);
        return ;
    }
    CipherUtils& utils = CipherUtils::getInstance();
    AppConfig& config = AppConfig::getInstance();

    std::string AESKey = config.runtimeValue("LocalAESKey");
    std::string AESIV = config.runtimeValue("LocalAESIV");

    auto res = utils.AES_Encrypt(file, size,
                                 (unsigned char*)AESKey.c_str(), (unsigned char*)AESIV.c_str());
    DEBUG_INFO("blockSize: " << res.second);
    ofs->write((char *)res.first.get(), res.second);
}

void FileManager::processQueue()
{
    while (true) {
        // get data
        std::unique_lock<std::mutex> lock(mtx);
        cv.wait(lock, [this]() { return !dataQueue.empty() || done; });

        if (done && dataQueue.empty()) {
            break;
        }

        nlohmann::json dataJson = dataQueue.front();
        dataQueue.pop();
        lock.unlock();
        handleFileBlock(dataJson);
    }
}

void FileManager::registerID(uint64_t id, const std::string &path,uint64_t fileSize)
{
    std::wstring_convert<std::codecvt_utf8<wchar_t>> converter;
    std::wstring filePath = converter.from_bytes(path);

    ID2FileInfo.emplace(id, std::make_pair(path, fileSize));
    std::shared_ptr<std::ofstream> ofs(new std::ofstream(filePath, std::ios::binary | std::ios::out));
    if (!ofs->is_open()) {
        DEBUG_INFO("文件打开失败"<<path);
    }
    ID2Ofstream[id] = ofs;
    ID2Lock[id] = std::make_shared<std::mutex>();
}

void FileManager::finishID(uint64_t id)
{
    auto lock = ID2Lock[id];
    std::unique_lock<std::mutex> guard(*lock, std::defer_lock);

    ID2FileInfo.erase(id);
    auto ofs = ID2Ofstream[id];
    ofs->close();
    ID2Ofstream.erase(id);

    ID2Lock.erase(id);
}

std::string FileManager::queryFilePath(uint64_t id)
{
    if (ID2FileInfo.count(id)) {
        return ID2FileInfo[id].first;
    }
    return "";
}

uint64_t FileManager::queryFileSize(uint64_t id)
{
    if (ID2FileInfo.count(id)) {
        return ID2FileInfo[id].second;
    }
    return 0;
}

void FileManager::commitToQueue(const nlohmann::json &data)
{
    {
        std::lock_guard<std::mutex> lock(mtx);
        dataQueue.push(data);
    }
    cv.notify_one();
}

void FileManager::setAutoWriteEncrypted(bool encrypted)
{
    this->encrypted = encrypted;
}

void FileManager::handleFileBlock(const nlohmann::json &dataJson)
{
    CipherUtils& utils = CipherUtils::getInstance();
    int blockOffset = dataJson["BlockOffset"];

    auto file = utils.Base64_Decode(dataJson["FileBase64"]);
    uint64_t transmissionID = dataJson["TransmissionID"];

    auto lock = ID2Lock[transmissionID];
    std::unique_lock<std::mutex> guard(*lock, std::defer_lock);
    // 当文件传输结束以后，则无法根据transmissionID找到对应的ofstream，此时直接结束
    if (!ID2Ofstream.count(transmissionID))
        return;
    auto ofs = ID2Ofstream[transmissionID];
    writeToDisk(ofs, blockOffset, (char*)file.first.get(), dataJson["BlockSize"], encrypted);
}
